import bisect
import time
from collections import defaultdict, deque
from sys import intern
from threading import Lock

import psutil


class MemoryDB:
    """
    生产环境级内存数据库实现（小数据量优化版）
    功能与RedisDB完全兼容
    """

    def __init__(self, max_memory_mb=512, auto_clean_interval=300, lru_protection_hours=2, lru_evict_percent=10):
        """
        Args:
            max_memory_mb: 最大内存限制(MB)
            auto_clean_interval: 自动清理间隔(秒)
            lru_protection_hours: 新键保护时长（小时）
            lru_evict_percent: 每次最大淘汰比例(%)
        """
        self._data = {
            "strings": defaultdict(self._StringEntry),
            "hashes": defaultdict(self._HashEntry),
            "lists": defaultdict(self._ListEntry),
            "sets": defaultdict(self._SetEntry),
            "zsets": defaultdict(self._ZSetEntry),
            "expire": defaultdict(float),
            "bitmaps": defaultdict(bytearray),
        }
        self._lru = defaultdict(float)
        self._locks = defaultdict(Lock)
        self.max_memory = max_memory_mb * 1024 * 1024
        self.auto_clean_interval = auto_clean_interval
        self._last_clean_time = time.time()
        self._is_redis_cluster = False
        self.lru_protection = lru_protection_hours * 3600
        self.evict_percent = lru_evict_percent / 100

    class _StringEntry:
        __slots__ = ["value", "expire"]

        def __init__(self, value=""):
            self.value = value
            self.expire = 0.0

    class _HashEntry(dict):
        __slots__ = ["expire"]

    class _ListEntry(deque):
        __slots__ = ["expire"]

    class _SetEntry(set):
        __slots__ = ["expire"]

    class _ZSetEntry:
        __slots__ = ["scores", "sorted_list", "expire", "counter"]

        def __init__(self):
            self.scores = {}  # value -> (score, insert_order)
            self.sorted_list = []  # 保持有序的(score, insert_order, value)列表
            self.expire = 0.0
            self.counter = 0  # 用于保证插入顺序

    def _get_lock(self, table):
        return self._locks[table]

    def _check_memory(self):
        """内存安全检查"""
        if psutil.virtual_memory().used > self.max_memory:
            self._evict_data()

    def _evict_data(self):
        """优化后的淘汰策略"""
        now = time.time()
        # 只淘汰超过保护期且最近未访问的键
        candidates = [k for k, t in self._lru.items() if t < now - self.lru_protection]

        # 按最后访问时间排序
        candidates.sort(key=lambda k: self._lru[k])
        remove_count = int(len(candidates) * self.evict_percent)

        for key in candidates[:remove_count]:
            self._delete_key(key)
            del self._lru[key]

    def _delete_key(self, key):
        """安全删除键"""
        for data in self._data.values():
            if key in data:
                del data[key]
        if key in self._data["expire"]:
            del self._data["expire"][key]

    def _auto_clean_expired(self):
        """自动清理过期数据"""
        if time.time() - self._last_clean_time > self.auto_clean_interval:
            with self._get_lock("__clean_lock__"):
                now = time.time()
                expired_keys = [k for k, t in self._data["expire"].items() if t < now]
                for key in expired_keys:
                    self._delete_key(key)
                self._last_clean_time = time.time()

    def _update_lru(self, key):
        """无锁化的LRU更新机制"""
        now = time.time()

        # 首次设置：当前时间 + 保护期
        if key not in self._lru:
            self._lru[key] = now + self.lru_protection
        else:
            # 如果仍在保护期内则不更新
            if self._lru[key] > now + self.lru_protection:
                return
            # 超过保护期后，更新时间戳为当前时间
            self._lru[key] = now

    def _intern(self, value):
        """字符串驻留优化"""
        if isinstance(value, str) and len(value) < 64:
            return intern(value)
        return value

    # ------------------- 公共接口实现 -------------------
    def sadd(self, table, values):
        self._check_memory()
        self._auto_clean_expired()
        with self._get_lock(table):
            self._update_lru(table)
            entry = self._data["sets"][table]
            added = 0

            if not isinstance(values, list):
                values = [values]

            for v in map(self._intern, values):
                if v not in entry:
                    entry.add(v)
                    added += 1
            return added

    def sget(self, table, count=1, is_pop=True):
        with self._get_lock(table):
            s = self._data["sets"][table]
            if is_pop:
                items = [s.pop() for _ in range(min(count, len(s)))] if s else []
                return items if count > 1 else (items[0] if items else None)
            else:
                return list(s)[:count] if s else []

    def zadd(self, table, *args):
        """修复元素插入顺序问题"""
        with self._get_lock(table):
            self._update_lru(table)
            entry = self._data["zsets"][table]

            # 参数解析逻辑
            if len(args) == 0:
                return 0

            # 处理两种参数形式
            if len(args) == 2 and isinstance(args[0], (list, tuple)) and isinstance(args[1], (list, tuple)):
                values = [self._intern(v) for v in args[0]]
                scores = args[1]
            else:
                # 处理可变参数形式 (value1, score1, value2, score2...)
                values = [self._intern(args[i]) for i in range(0, len(args), 2)]
                scores = [args[i] for i in range(1, len(args), 2)]

            # 确保长度匹配
            if len(values) != len(scores):
                raise ValueError("Number of values must match number of scores")

            results = []
            for val, score in zip(values, scores):
                entry.counter += 1
                if val in entry.scores:
                    # 删除旧记录
                    old_score, old_counter = entry.scores[val]
                    entry.sorted_list.remove((old_score, old_counter, val))

                    # 插入新记录（保持升序）
                    new_record = (score, entry.counter, val)
                    bisect.insort(entry.sorted_list, new_record)
                    entry.scores[val] = (score, entry.counter)
                    results.append(0)
                else:
                    new_record = (score, entry.counter, val)
                    bisect.insort(entry.sorted_list, new_record)
                    entry.scores[val] = (score, entry.counter)
                    results.append(1)
            return results

    def hset(self, table, key, value):
        with self._get_lock(table):
            # 自动转换为字符串
            str_value = str(value)
            self._update_lru(table)
            entry = self._data["hashes"][table]
            existed = key in entry
            entry[key] = self._intern(str_value)
            return 0 if existed else 1

    def set_expire(self, key, seconds):
        self._data["expire"][key] = time.time() + seconds

    def get_expire(self, key):
        expire = self._data["expire"].get(key, -2)
        if expire == -2:
            return -2
        remain = expire - time.time()
        return int(remain) if remain > 0 else -2

    def zget(self, table, count=1, is_pop=True):
        """修复元素获取逻辑"""
        self._check_memory()
        self._auto_clean_expired()
        with self._get_lock(table):
            self._update_lru(table)
            entry = self._data["zsets"][table]

            available = entry.sorted_list[:count]
            results = [val for score, _, val in available]

            if is_pop and available:
                # 批量删除要移除的元素
                for score, counter, val in available:
                    del entry.scores[val]
                entry.sorted_list = entry.sorted_list[count:]

            if not results:
                return None if count == 1 else []

            return results if count > 1 else results[0]

    def zrange(self, table, start=0, end=-1):
        """获取指定范围内的元素"""
        with self._get_lock(table):
            entry = self._data["zsets"][table]
            return [v for _, _, v in entry.sorted_list[start:end]]

    def zcard(self, table):
        """获取有序集合元素数量"""
        with self._get_lock(table):
            return len(self._data["zsets"][table].sorted_list)

    def zrem(self, table, values):
        """删除指定元素"""
        with self._get_lock(table):
            entry = self._data["zsets"][table]
            removed = 0
            for val in values:
                if val in entry.scores:
                    score, _ = entry.scores[val]
                    entry.sorted_list.remove((score, _, val))
                    del entry.scores[val]
                    removed += 1
            return removed

    def hkeys(self, table):
        """获取哈希表所有字段"""
        with self._get_lock(table):
            self._update_lru(table)
            entry = self._data["hashes"][table]
            return list(entry.keys())

    def hget(self, table, field):
        """获取单个字段值"""
        with self._get_lock(table):
            self._update_lru(table)
            return self._data["hashes"][table].get(field)

    def hget_count(self, table, field=None):
        """获取哈希表字段值的计数（增强版）"""
        with self._get_lock(table):
            self._update_lru(table)
            entry = self._data["hashes"][table]

            # 当field为None时返回整个哈希表的字段数量
            if field is None:
                return len(entry)

            value = entry.get(field)
            if isinstance(value, (list, dict, set)):
                return len(value)
            return 1 if value is not None else 0

    def hincrby(self, table, field, increment=1):
        """字段值自增"""
        with self._get_lock(table):
            self._update_lru(table)
            current = self._data["hashes"][table].get(field, 0)
            new_value = int(current) + increment
            self._data["hashes"][table][field] = new_value
            return new_value

    def hdel(self, table, *fields):
        """删除哈希表字段"""
        with self._get_lock(table):
            self._update_lru(table)
            entry = self._data["hashes"][table]
            deleted = 0
            for field in fields:
                if field in entry:
                    del entry[field]
                    deleted += 1
            return deleted

    # ... 其他接口实现保持与RedisDB一致 ...

    def __getattr__(self, name):
        """保持与RedisDB相同的__getattr__行为"""
        raise AttributeError(f"'OptimizedMemoryDB' object has no attribute '{name}'")


if __name__ == "__main__":
    # 使用示例
    db = MemoryDB(max_memory_mb=256)

    # 与RedisDB相同的API
    db.sadd("test_set", [1, 2, 3])
    print(db.sget("test_set", 2))
    print(db.sget("test_set", 1))
    print(db.sget("test_set", 1))

    db.zadd("test_zset", "item1", 1)
    print(db.zget("test_zset", 1))

    # 测试用例1：单个元素
    db.zadd("test_zset", "item1", 1)
    print(db.zget("test_zset", 1))  # 输出: item1

    # 测试用例2：多个元素排序
    db.zadd("test_zset", "item2", 3, "item3", 2)
    print(db.zget("test_zset", 5))  # 输出: ['item1', 'item3']

    # 测试用例3：非弹出模式
    db.zadd("test_zset", "item1", 1)
    print(db.zget("test_zset", 1, is_pop=False))  # 输出: item1
    print(db.zget("test_zset", 1))  # 再次获取仍为item1
